/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Definitions of the member functions for the particle_set
    classes. \ingroup sfrhist */

// $Id$

#include "particle_set.h" 

#include "boost/lambda/lambda.hpp"
#include "blitz-fits.h"
#include "model.h"

using namespace std;
using namespace boost::lambda;

/** Finds a particle with the specified ID. */
int mcrx::particle_set_generic::find_ID(unsigned int id) const {
  // if the particles are sorted in id order, we can use binary search
  // on the vector instead of the map search. in fact we don't need
  // the map at all then
  const std::map<unsigned int,int>::const_iterator i = order_.find(id);
  if (i == order_.end()) throw illegal_id(); 
  return i->second;
};


void 
mcrx::particle_set_generic::convert_units(T_float mass, T_float length, 
					  T_float time)
{
  transform (m_.begin(), m_.end(), m_.begin(), _1 * mass);
  transform (pos_.begin(), pos_.end(), pos_.begin(), 
	     ret<vec3d>(_1 * length));
  transform (vel_.begin(), vel_.end(), vel_.begin(),    
	     ret<vec3d>(_1 * (length/time) ));
  transform (h_.begin(), h_.end(), h_.begin(), _1 * length);
}

mcrx::particle_set_generic&
mcrx::particle_set_generic::operator+=(const particle_set_generic& rhs)
{
  const int Nold=n();
  copy(rhs.id_.begin(), rhs.id_.end(), back_inserter(id_));
  copy(rhs.m_.begin(), rhs.m_.end(), back_inserter(m_));
  copy(rhs.pos_.begin(), rhs.pos_.end(), back_inserter(pos_));
  copy(rhs.vel_.begin(), rhs.vel_.end(), back_inserter(vel_));
  copy(rhs.h_.begin(), rhs.h_.end(), back_inserter(h_));
  // we must reconstruct the order map to point to the new position
  for(map<unsigned int, int>::const_iterator i=rhs.order_.begin(); 
      i!=rhs.order_.end(); ++i) {
    const int id=i->first;
    assert(order_.find(id)==order_.end());
    order_[i->first]=i->second+Nold;
  }
  return *this;
}

void 
mcrx::gas_particle_set::convert_units(T_float mass, T_float length, 
				      T_float time)
{
  particle_set_generic::convert_units(mass, length, time);

  transform (sfr_.begin(), sfr_.end(), sfr_.begin(), _1 * (mass/time) );
  
  // metallicity is a fraction, unitless
  // temp is in K, so is not affected by these unit changes
}


mcrx::gas_particle_set&
mcrx::gas_particle_set::operator+=(const gas_particle_set& rhs)
{
  particle_set_generic::operator+=(rhs);
  copy(rhs.z_.begin(), rhs.z_.end(), back_inserter(z_));
  copy(rhs.sfr_.begin(), rhs.sfr_.end(), back_inserter(sfr_));
  copy(rhs.temp_.begin(), rhs.temp_.end(), back_inserter(temp_));
  copy(rhs.teff_.begin(), rhs.teff_.end(), back_inserter(teff_));
  copy(rhs.pid_.begin(), rhs.pid_.end(), back_inserter(pid_));
  return *this;
}


void 
mcrx::star_particle_set::convert_units(T_float mass, T_float length, 
				       T_float time)
{
  particle_set_generic::convert_units(mass, length, time);

  transform (tf_.begin(), tf_.end(), tf_.begin(), _1 * time );
  transform (mcr_.begin(), mcr_.end(), mcr_.begin(), _1 * mass );

  // metallicity is a fraction, unitless
  // parent ID is unitless
}


mcrx::star_particle_set&
mcrx::star_particle_set::operator+=(const star_particle_set& rhs)
{
  particle_set_generic::operator+=(rhs);
  copy(rhs.z_.begin(), rhs.z_.end(), back_inserter(z_));
  copy(rhs.tf_.begin(), rhs.tf_.end(), back_inserter(tf_));
  copy(rhs.pid_.begin(), rhs.pid_.end(), back_inserter(pid_));
  copy(rhs.mcr_.begin(), rhs.mcr_.end(), back_inserter(mcr_));
  return *this;
}


void 
mcrx::dm_particle_set::convert_units(T_float mass, T_float length, 
				       T_float time)
{
  particle_set_generic::convert_units(mass, length, time);
}


mcrx::dm_particle_set&
mcrx::dm_particle_set::operator+=(const dm_particle_set& rhs)
{
  particle_set_generic::operator+=(rhs);
  return *this;
}


void 
mcrx::bh_particle_set::convert_units(T_float mass, T_float length, 
				       T_float time)
{
  particle_set_generic::convert_units(mass, length, time);

  transform (mbh_.begin(), mbh_.end(), mbh_.begin(), _1 * mass );
  transform (mdotbh_.begin(), mdotbh_.end(), mdotbh_.begin(),
	     _1 * (mass/time) );
}


mcrx::bh_particle_set&
mcrx::bh_particle_set::operator+=(const bh_particle_set& rhs)
{
  particle_set_generic::operator+=(rhs);
  copy(rhs.mbh_.begin(), rhs.mbh_.end(), back_inserter(mbh_));
  copy(rhs.mdotbh_.begin(), rhs.mdotbh_.end(), back_inserter(mdotbh_));
  return *this;
}


/** Writes a chunk of the particle data to a FITS table. */
void
mcrx::particle_set_generic::write_table_chunk (int start_row,
					       int start_particle,
					       int n_chunk,
					       CCfits::ExtHDU& hdu) const
{
  CCfits::Column& c_id = hdu.column("ID") ;
  CCfits::Column& c_m = hdu.column("mass") ;
  CCfits::Column& c_pos = hdu.column("position") ;
  CCfits::Column& c_vel = hdu.column("velocity") ;
  CCfits::Column& c_h = hdu.column("radius") ;

  make_onedva(id_).write_column(hdu, c_id, start_row, start_particle, n_chunk);
  make_onedva(m_).write_column(hdu, c_m, start_row, start_particle, n_chunk);
  
  write(c_pos, pos_, start_row, start_particle, n_chunk);
  write(c_vel, vel_, start_row, start_particle, n_chunk);
  make_onedva(h_).write_column(hdu, c_h, start_row, start_particle, n_chunk);
}

void mcrx::particle_set_generic::write_ascii_table(ostream& os) const
{
  for(int i=0; i<n(); ++i) {
    write_ascii_table_line(os, i);
    os << "\n";
  } 
    os << "\n\n\n";
}

void mcrx::particle_set_generic::write_ascii_table_line(ostream& os,
							int i) const
{
  os << id_[i] << '\t' << m_[i] << '\t' 
     << pos_[i][0] << '\t' << pos_[i][1] << '\t' << pos_[i][2] << '\t'
     << vel_[i][0] << '\t' << vel_[i][1] << '\t' << vel_[i][2] << '\t'
     << h_[i] << '\t';
}

void
mcrx::particle_set_generic::load_particledata (CCfits::ExtHDU& hdu)
{
  CCfits::Column& c_id = hdu.column("ID") ;
  CCfits::Column& c_m = hdu.column("mass") ;
  CCfits::Column& c_pos = hdu.column("position") ;
  CCfits::Column& c_vel = hdu.column("velocity") ;
  CCfits::Column& c_h = hdu.column("radius") ;

  const int n= c_id.rows();

  c_id.read(id_, 1, n );
  c_m.read(m_, 1, n);
  read (c_pos, pos_);
  read (c_vel, vel_);
  c_h.read(h_, 1, n);
}

void
mcrx::gas_particle_set::write_table_chunk (int start_row,
					   int start_particle,
					   int n_chunk,
					   CCfits::ExtHDU& hdu) const
{
  particle_set_generic::write_table_chunk(start_row, start_particle, n_chunk,
					  hdu);

  CCfits::Column& c_z = hdu.column("metallicity") ;
  CCfits::Column& c_sfr = hdu.column("sfr") ;
  CCfits::Column& c_temp = hdu.column("temperature") ;
  CCfits::Column& c_teff = hdu.column("teff") ;

  make_onedva(z_).write_column(hdu, c_z, start_row, start_particle, n_chunk);
  make_onedva(sfr_).write_column(hdu, c_sfr, start_row, start_particle, n_chunk);
  make_onedva(temp_).write_column(hdu, c_temp, 
				  start_row, start_particle, n_chunk);
  make_onedva(teff_).write_column(hdu, c_teff,
				  start_row, start_particle, n_chunk);
}

void mcrx::gas_particle_set::write_ascii_table_line(ostream& os,
						    int i) const
{
  particle_set_generic::write_ascii_table_line(os, i);
  os << z_[i] << '\t' << sfr_[i] << '\t' 
     << temp_[i] << '\t' << teff_[i] << '\t';
}

void mcrx::gas_particle_set::load_particledata (CCfits::ExtHDU& hdu)
{
  particle_set_generic::load_particledata (hdu);

  CCfits::Column& c_z = hdu.column("metallicity") ;
  CCfits::Column& c_sfr = hdu.column("sfr") ;
  CCfits::Column& c_temp = hdu.column("temperature") ;
  CCfits::Column& c_teff = hdu.column("teff") ;

  const int n= c_z.rows();

  c_z.read(z_, 1, n );
  c_sfr.read(sfr_, 1, n);
  c_temp.read(temp_, 1, n);
  c_teff.read(teff_, 1, n);
}


void
mcrx::star_particle_set::write_table_chunk (int start_row,
					    int start_particle,
					    int n_chunk,
					    CCfits::ExtHDU& hdu) const
{
  particle_set_generic::write_table_chunk(start_row, start_particle, 
					  n_chunk, hdu);

  CCfits::Column& c_z = hdu.column("metallicity") ;
  CCfits::Column& c_tf = hdu.column("formation_time") ;
  CCfits::Column& c_pid = hdu.column("parent_ID") ;
  CCfits::Column& c_mcr = hdu.column("creation_mass") ;

  make_onedva(z_).write_column(hdu, c_z, start_row, start_particle, n_chunk);
  // The presence of the tf field is a little suboptimal here, we want
  // age. What we do is let the set write the tf and then add the age
  // along with the SED in sfrhist_snashot.
  make_onedva(tf_).write_column(hdu, c_tf, start_row, start_particle, n_chunk);
  make_onedva(pid_).write_column(hdu, c_pid, 
				 start_row, start_particle, n_chunk);
  make_onedva(mcr_).write_column(hdu, c_mcr,
				 start_row, start_particle, n_chunk);
}
  
void mcrx::star_particle_set::write_ascii_table_line(ostream& os,
						     int i) const
{
  particle_set_generic::write_ascii_table_line(os, i);
  os << z_[i] << '\t' << tf_[i] << '\t' 
     << pid_[i] << '\t' << mcr_[i] << '\t';
}

void mcrx::star_particle_set::load_particledata (CCfits::ExtHDU& hdu)
{
  particle_set_generic::load_particledata (hdu);

  CCfits::Column& c_z = hdu.column("metallicity") ;
  CCfits::Column& c_tf = hdu.column("formation_time") ;
  CCfits::Column& c_pid = hdu.column("parent_ID") ;
  CCfits::Column& c_mcr = hdu.column("creation_mass") ;

  const int n= c_z.rows();

  c_z.read(z_, 1, n );
  c_tf.read(tf_, 1, n);
  c_pid.read(pid_, 1, n);
  c_mcr.read(mcr_, 1, n);
}

void
mcrx::dm_particle_set::write_table_chunk (int start_row,
					  int start_particle,
					  int n_chunk,
					  CCfits::ExtHDU& hdu) const
{
  particle_set_generic::write_table_chunk(start_row, start_particle, 
					  n_chunk, hdu);
}


void mcrx::dm_particle_set::load_particledata (CCfits::ExtHDU& hdu)
{
  particle_set_generic::load_particledata (hdu);
}


void
mcrx::bh_particle_set::write_table_chunk (int start_row,
					  int start_particle,
					  int n_chunk,
					  CCfits::ExtHDU& hdu) const
{
  particle_set_generic::write_table_chunk(start_row, start_particle, 
					  n_chunk, hdu);
  // if the columns exist, write them
  try {
    CCfits::Column& c_mbh = hdu.column("bh_mass") ;
    make_onedva(mbh_).write_column(hdu, c_mbh, start_row, 
				   start_particle, n_chunk);
  }
  catch (CCfits::Table::NoSuchColumn&) {}

  try {
    CCfits::Column& c_mdotbh = hdu.column("bh_accretion_rate") ;
    make_onedva(mdotbh_).write_column(hdu, c_mdotbh, start_row,
				      start_particle, n_chunk);
  }
  catch (CCfits::Table::NoSuchColumn&) {}

}

void mcrx::bh_particle_set::load_particledata (CCfits::ExtHDU& hdu)
{
  particle_set_generic::load_particledata (hdu);

  CCfits::Column& c_mbh = hdu.column("bh_mass") ;
  CCfits::Column& c_mdotbh = hdu.column("bh_accretion_rate") ;

  const int n= c_mbh.rows();

  c_mbh.read(mbh_, 1, n );
  c_mdotbh.read(mdotbh_, 1, n);
}



void mcrx::star_particle_set::calculate_creation_mass(T_float time,
						      const mcrx::stellarmodel& m)
{
  for (int i=0; i<n(); ++i) {
    const T_float age = time - tf_[i];
    const T_float z = z_[i];
    const T_float f = m.get_mass_fraction(age, z);
    mcr_[i] = m_[i]/f;
  }
}

